//
// Created by ECNU on 2024/2/22.
//
#include <nanofaiss/impl/ResultHandler.h>
#include <nanofaiss/utils/distances.h>

#include <algorithm>
#include <omp.h>
#include <memory>


extern "C" {

#undef FINTEGER
#define FINTEGER long

/* declare BLAS functions, see http://www.netlib.org/clapack/cblas/ */
int sgemm_(
        const char* transa,
        const char* transb,
        FINTEGER* m,
        FINTEGER* n,
        FINTEGER* k,
        const float* alpha,
        const float* a,
        FINTEGER* lda,
        const float* b,
        FINTEGER* ldb,
        float* beta,
        float* c,
        FINTEGER* ldc);
}

namespace faiss {

int distance_compute_min_k_reservoir = 1000; //TODO: Need to be fix to 100
int distance_compute_blas_threshold = 20;
int distance_compute_blas_query_bs = 4096;
int distance_compute_blas_database_bs = 1024;


/// 采用inner product的方式计算L2范数
template<class BlockResultHandler>
void exhaustive_L2sqr_blas_default_impl(
        const float* x,
        const float* y,
        size_t d,
        size_t nx,
        size_t ny,
        BlockResultHandler& res,
        const float *y_norms = nullptr) {
    if (nx == 0 || ny == 0)
        return;

    const size_t bs_x = distance_compute_blas_query_bs;
    const size_t bs_y = distance_compute_blas_database_bs;

    // compute method: |x - y|^2 = x^2 + y^2 - 2xy
    std::unique_ptr<float[]> ip_block(new float[bs_x * bs_y]); // float[]表示float类型数组, 那么就可以使用array的方法
    std::unique_ptr<float[]> x_norms(new float[nx]);
    std::unique_ptr<float[]> del2;

    // compute the norm of x
    fvec_norms_L2sqr(x_norms.get(), x, d, nx);

    if (!y_norms) {
        float *y_norms2 = new float[ny];
        del2.reset(y_norms2);
        fvec_norms_L2sqr(y_norms2, y, d, ny);
        y_norms = y_norms2;
    }

    // batch compute the result
    // loop for batch x
    for (size_t i0 = 0; i0 < nx; i0 += bs_x) {
        size_t i1 = i0 + bs_x;
        i1 = std::min(i1, nx);

        res.begin_multiple(i0, i1);

        // loop for batch y
        for (size_t j0 = 0; j0 < ny; j0 += bs_y) {
            size_t j1 = j0 + bs_y;
            j1 = std::min(j1, ny);
            // compute inner product
            {
                float one = 1, zero = 0;
                FINTEGER nyi = j1 - j0, nxi = i1 - i0, di = d;
                sgemm_("Transpose",
                       "Not transpose",
                       &nyi,
                       &nxi,
                       &di,
                       &one,
                       y + j0 * d,
                       &di,
                       x + i0 * d,
                       &di,
                       &zero,
                       ip_block.get(),
                       &nyi);
            }

            // inner loop calculate distance
#pragma omp parallel for
            for (size_t i = i0; i != i1; i++) {
                float *ip_line = ip_block.get() + (i - i0) * (j1 - j0);

                for (size_t j = j0; j < j1; j++) {
                    float ip = *ip_line;
                    float dis = x_norms[i] + y_norms[j] - 2 * ip;
                    if (dis < 0)
                        dis = 0;

                    *ip_line = dis; // direct save the result in ip_line -> ip_block
                    ip_line++;
                }
            }
            res.add_results(j0, j1, ip_block.get());
        }
        res.end_multiple();
    }
}



template<class BlockResultHandler>
void exhaustive_L2sqr_blas(
        const float *x,
        const float *y,
        size_t d,
        size_t nx,
        size_t ny,
        BlockResultHandler &res,
        const float *y_norms = nullptr) {
    exhaustive_L2sqr_blas_default_impl(x, y, d, nx, ny, res, y_norms);

}

template<class BlockResultHandler>
void exhaustive_L2sqr_seq(
        const float *x,
        const float *y,
        size_t d,
        size_t nx,
        size_t ny,
        BlockResultHandler &res) {
    using SingleResultHandler =
            typename BlockResultHandler::SingleResultHandler;
    int nt = std::min(static_cast<int>(nx), omp_get_max_threads());

#pragma omp parallel num_threads(nt)
    {
        SingleResultHandler resi(res);
#pragma omp for
        for (size_t i = 0; i != nx; i++) {
            const float *xi = x + d * i;
            resi.begin(i);
            for (size_t j = 0; j != ny; j++) {
                const float *yj = y + j * d;
                float disij = fvec_L2sqr(xi, yj, d);
                resi.add_result(disij, j);
            }
            resi.end();
        }

    }



}

template<class BlockResultHandler>
void knn_L2sqr_select(
        const float *x,
        const float *y,
        size_t d,
        size_t nx,
        size_t ny,
        BlockResultHandler& res,
        const float *y_norm2) {
    if (nx < distance_compute_blas_threshold) {
        exhaustive_L2sqr_seq(x, y, d, nx, ny, res);
    } else {
//        exhaustive_L2sqr_blas(x, y, d, nx, ny, res, y_norm2);
        exhaustive_L2sqr_seq(x, y, d, nx, ny, res); //TODO: Need to recover
    }
}

void knn_L2sqr(
        const float* x,
        const float* y,
        size_t d,
        size_t nx,
        size_t ny,
        size_t k,
        float* distances,
        int64_t* ids,
        const float* y_norm2) {
    // actually there are 3 methods depends on the number of k
    // Top1 optimization | HeapHandle | Reservoir
    // TODO: Now we just use HeapHandle
    if (k == -1) { //TODO: Need to fix to 1

    } else if (k < distance_compute_min_k_reservoir) {
        HeapBlockResultHandler<CMax<float, int64_t>> res(nx, distances, ids, k);
        knn_L2sqr_select(x, y, d, nx, ny, res, y_norm2);
    } else {

    }

}


void knn_L2sqr(
        const float *x,
        const float *y,
        size_t d,
        size_t nx,
        size_t ny,
        float_query_t *res,
        const float *y_norm) {
    FAISS_THROW_IF_NOT(nx == res->nq);
    knn_L2sqr(x, y, d, nx, ny, res->k,res->distances, res->ids, y_norm);
}



void fvec_norms_L2sqr(
        float *__restrict nr,
        const float* __restrict x,
        size_t d,
        size_t nx) {
#pragma omp parallel for if (nx > 10000)
    for (int64_t i = 0; i < nx; i++)
        nr[i] = fvec_norm_L2sqr(x + i * d, d);
}

}
